import websocket
import requests
import uuid
import json
import urllib.parse
import os
import time
import logging
from typing import Optional, Dict, Any, List

# 设置日志记录
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
)


class ComfyUIError(Exception):
    """自定义 ComfyUI 客户端异常基类"""

    pass


class WorkflowError(ComfyUIError):
    """工作流相关的错误"""

    pass


class ComfyUIClient:
    """
    一个用于与 ComfyUI API 交互的客户端。

    该类封装了与 ComfyUI 后端通信的通用逻辑，包括：
    - 连接到服务器
    - 上传图片
    - 加载和准备工作流
    - 提交任务并监控执行
    - 下载生成的图片

    通过定义“节点映射”，用户可以轻松地将输入参数与任何 ComfyUI 工作流中的特定节点进行关联，
    从而实现对任意工作流的动态控制，而无需修改代码。
    """

    def __init__(self, server_address: str,  client_id: Optional[str] = None):
        """
        初始化 ComfyUI 客户端。

        Args:
            server_address (str): ComfyUI 服务器的 IP 地址或域名。
            port (int): ComfyUI 服务器的端口。
            client_id (Optional[str]): 客户端 ID。如果未提供，将自动生成一个。
        """
        if not server_address:
            raise ValueError("服务器地址不能为空。")

        self.server_address = server_address
        self.base_url = f"https://{self.server_address}"
        self.ws_url = f"wss://{self.server_address}/ws"
        self.client_id = client_id or uuid.uuid4().hex
        logging.info(
            f"ComfyUI 客户端已初始化。服务器: {self.base_url}, 客户端ID: {self.client_id}"
        )

    def _get_url(self, endpoint: str) -> str:
        """构建完整的 API URL"""
        return f"{self.base_url}{endpoint}"

    def _send_request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]:
        """发送 HTTP 请求的通用方法"""
        url = self._get_url(endpoint)
        try:
            response = requests.request(method, url, **kwargs)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.HTTPError as http_err:
            logging.error(f"HTTP 错误: {http_err} - {response.text}")
            raise ComfyUIError(f"HTTP 请求失败: {http_err}") from http_err
        except requests.exceptions.RequestException as req_err:
            logging.error(f"请求异常: {req_err}")
            raise ComfyUIError(f"无法连接到 ComfyUI 服务器: {req_err}") from req_err

    def upload_image(
        self, image_path: str, subfolder: str = "", image_type: str = "input"
    ) -> Dict[str, Any]:
        """
        上传图片到 ComfyUI 服务器。

        Args:
            image_path (str): 本地图片文件的路径。
            subfolder (str, optional): 在 ComfyUI input 目录下的子文件夹。默认为 ""。
            image_type (str, optional): 上传类型，通常是 "input"。默认为 "input"。

        Returns:
            Dict[str, Any]: 包含上传图片信息的字典。

        Raises:
            FileNotFoundError: 如果图片路径不存在。
            ComfyUIError: 如果上传失败。
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图片文件未找到: {image_path}")

        filename = os.path.basename(image_path)
        with open(image_path, "rb") as f:
            files = {"image": (filename, f)}
            data = {"overwrite": "true", "subfolder": subfolder, "type": image_type}
            logging.info(f"正在上传图片 '{filename}' 到子目录 '{subfolder}'...")
            result = self._send_request("POST", "/upload/image", files=files, data=data)
            logging.info("图片上传成功。")
            return result

    def _queue_prompt(self, prompt_workflow: Dict[str, Any]) -> str:
        """
        将准备好的工作流加入队列。

        Args:
            prompt_workflow (Dict[str, Any]): API 格式的工作流。

        Returns:
            str: 任务的 Prompt ID。
        """
        data = {"prompt": prompt_workflow, "client_id": self.client_id}
        response = self._send_request("POST", "/prompt", json=data)
        if "prompt_id" not in response:
            raise ComfyUIError(f"提交任务失败，响应中缺少 'prompt_id': {response}")
        prompt_id = response["prompt_id"]
        logging.info(f"任务已提交，Prompt ID: {prompt_id}")
        return prompt_id

    def _wait_for_completion(self, prompt_id: str):
        """通过 WebSocket 等待任务执行完成。"""
        ws = websocket.WebSocket()
        try:
            ws.connect(f"{self.ws_url}?clientId={self.client_id}")
            logging.info("WebSocket 已连接，正在等待任务执行...")
            while True:
                out = ws.recv()
                if isinstance(out, str):
                    message = json.loads(out)
                    if message["type"] == "executing":
                        data = message["data"]
                        # 当 node 为 null 时，表示执行结束
                        if data["node"] is None and data.get("prompt_id") == prompt_id:
                            logging.info(f"Prompt {prompt_id} 执行完成。")
                            break
        except websocket.WebSocketConnectionClosedException:
            logging.warning("WebSocket 连接已关闭。")
        except Exception as e:
            raise ComfyUIError(f"WebSocket 监听出错: {e}") from e
        finally:
            if ws.connected:
                ws.close()

    def _get_history(self, prompt_id: str) -> Dict[str, Any]:
        """获取指定任务的历史记录。"""
        return self._send_request("GET", f"/history/{prompt_id}")

    def _get_images(
        self, prompt_id: str, save_node_id: str, output_dir: str
    ) -> List[str]:
        """
        从历史记录中下载由特定 SaveImage 节点生成的图片。

        Args:
            prompt_id (str): 任务的 Prompt ID。
            save_node_id (str): 工作流中 "SaveImage" 节点的 ID。
            output_dir (str): 保存图片的本地目录。

        Returns:
            List[str]: 下载到本地的图片路径列表。
        """
        os.makedirs(output_dir, exist_ok=True)
        history = self._get_history(prompt_id).get(prompt_id)
        print("历史记录内容：", json.dumps(history, indent=2, ensure_ascii=False))
        if not history:
            raise WorkflowError(f"在历史记录中找不到 Prompt ID: {prompt_id}")

        outputs = history.get("outputs", {})
        if save_node_id not in outputs:
            raise WorkflowError(f"在历史输出中找不到 SaveImage 节点 ID: {save_node_id}")

        images_output = outputs[save_node_id].get("images", [])
        downloaded_paths = []

        for img_info in images_output:
            filename = img_info["filename"]
            subfolder = img_info.get("subfolder", "")
            img_type = img_info.get("type", "output")

            image_url = f"{self.base_url}/view?filename={urllib.parse.quote(filename)}&subfolder={urllib.parse.quote(subfolder)}&type={img_type}"
            logging.info(f"正在下载图片: {image_url}")

            try:
                img_response = requests.get(image_url, stream=True)
                img_response.raise_for_status()

                local_path = os.path.join(output_dir, filename)
                with open(local_path, "wb") as f:
                    for chunk in img_response.iter_content(chunk_size=8192):
                        f.write(chunk)
                downloaded_paths.append(local_path)
                logging.info(f"图片已保存到: {local_path}")
            except requests.exceptions.RequestException as e:
                logging.error(f"下载图片失败: {e}")
                # 继续尝试下载其他图片

        return downloaded_paths

    def run_workflow(
        self,
        workflow_api_json_path: str,
        node_mapping: Dict[str, Dict[str, Any]],
        inputs: Dict[str, Any],
        output_dir: str = "generated_images",
    ) -> List[str]:
        """
        执行一个 ComfyUI 工作流。
        这是最核心的通用方法，所有具体任务（如文生图、图生图）都基于此。

        Args:
            workflow_api_json_path (str): API 格式的工作流 JSON 文件路径。
            node_mapping (Dict[str, Dict[str, Any]]):
                一个定义了如何将 `inputs` 映射到工作流节点的字典。
                格式: {"<input_name>": {"node_id": "<id>", "field": "<field_name>"}, ...}
                例如: {"positive_prompt": {"node_id": "6", "field": "text"}}
            inputs (Dict[str, Any]): 包含要更新到工作流中的参数的字典。
                                     键名应与 `node_mapping` 中的键名对应。
            output_dir (str, optional): 保存输出图片的目录。默认为 "generated_images"。

        Returns:
            List[str]: 生成的图片在本地的路径列表。
        """
        try:
            # 1. 加载工作流模板
            with open(workflow_api_json_path, "r", encoding="utf-8") as f:
                prompt_workflow = json.load(f)
        except FileNotFoundError:
            raise WorkflowError(f"工作流文件未找到: {workflow_api_json_path}")
        except json.JSONDecodeError:
            raise WorkflowError(f"无法解析工作流文件: {workflow_api_json_path}")

        # 2. 准备工作流：上传必要的图片并更新节点输入
        for input_name, value in inputs.items():
            if input_name not in node_mapping:
                logging.warning(f"输入 '{input_name}' 在节点映射中未定义，将被忽略。")
                continue

            mapping_info = node_mapping[input_name]
            node_id = mapping_info["node_id"]
            field = mapping_info["field"]

            if node_id not in prompt_workflow:
                raise WorkflowError(f"节点 ID '{node_id}' 在工作流中不存在。")

            # 特殊处理：如果输入是图片路径，则先上传
            if field == "image" and isinstance(value, str) and os.path.exists(value):
                logging.info(f"检测到图片输入 '{input_name}'，路径: {value}")
                upload_info = self.upload_image(value)
                prompt_workflow[node_id]["inputs"][field] = upload_info["name"]
            else:
                prompt_workflow[node_id]["inputs"][field] = value

        # 3. 提交任务并等待完成
        prompt_id = self._queue_prompt(prompt_workflow)
        self._wait_for_completion(prompt_id)

        # 4. 下载结果
        save_node_mapping = node_mapping.get("save_image")
        if not save_node_mapping:
            raise WorkflowError("节点映射中必须包含 'save_image' 以指定输出节点。")

        save_node_id = save_node_mapping["node_id"]
        return self._get_images(prompt_id, save_node_id, output_dir)


if __name__ == "__main__":
    # --- 使用示例 ---
    # 你的同事们未来只需要这样使用：

    # 1. 初始化客户端
    # 请确保 ComfyUI 服务器正在运行
    client = ComfyUIClient("f8de59b5a3c2.ngrok-free.app")

    # ======================================================================
    # 示例1: 运行一个基本的“文生图” (Text-to-Image)
    # ======================================================================
    logging.info("\n--- 示例: 运行多图融合任务 ---")

    # 定义这个工作流中，我们关心的、需要动态修改的节点
    # 这是关键！你的同事只需要根据他的工作流JSON文件填写一次
    ttp_node_mapping = {
        "background_image": {"node_id": "54", "field": "image"},
        "object_1_image": {"node_id": "9", "field": "image"},
        "object_2_image": {"node_id": "52", "field": "image"},
        "prompt": {"node_id": "16", "field": "text"},
        "unet_loader": {"node_id": "1", "field": "unet_name"},
        "sampler_seed": {"node_id": "13", "field": "seed"},
        "sampler_denoise": {"node_id": "13", "field": "denoise"},
        "save_image": {
            "node_id": "67",
            "field": "filename_prefix",
        },  # 我们需要知道哪个是保存图片的节点
    }

    # 定义要输入的参数
    ttp_inputs = {
        "background_image": "../static/background.png",
        "object_1_image": "../static/object1.png",
        "object_2_image": "../static/object2.png",
        "prompt": "让图中女人拿着一瓶啤酒坐在沙发上，拉近拍摄距离，保持角色不变，调整为正确的人体比例",
        "unet_loader": "flux1-kontext-dev.safetensors",  # 确保这个模型存在
        "sampler_seed": int(time.time()),  # 使用随机种子
        "sampler_denoise": 1.0,
    }

    try:
        # 假设你的文生图工作流文件叫 'TTP.json'
        # 你可以从你的代码中导出这个json
        generated_images_ttp = client.run_workflow(
            workflow_api_json_path="../static/workflows/multi_image_fusion_3.0 .json",  # 使用你的文生图工作流
            node_mapping=ttp_node_mapping,
            inputs=ttp_inputs,
            output_dir="generated_images/multi_imgs",
        )
        logging.info(f"多图融合成功！图片保存在: {generated_images_ttp}")
    except (ComfyUIError, FileNotFoundError) as e:
        logging.error(f"多图融合任务失败: {e}")